﻿typedef struct
{
	float4 position;
	float4 color;
} Light;		//32 bytes

typedef struct
{	
	union
	{
		struct {	//sphere
			float4 position;		//pos 0
			float radius;			//pos 16
		} sphere;
		struct {	//plane
			float4 normal;			//pos 0
			float offset;			//pos 16			
		} plane;
		struct {	//triangle
			float4 p1;				//pos 0
			float4 p2;				//pos 16
			float4 p3;				//pos 32
		} triangle;
	} properties;

	uchar type;		//0 = sphere, 1 = plane, 2 = triangle		pos 48
	uchar surface;											  //pos 49
} Object;		//64 bytes

typedef struct
{
	union
	{
		uint flags;
		float split;		//if node
		uint numObjects;	//if leaf
	} first;
	union
	{
		uint aboveChild;	//if node
		uint objectIndex;	//if leaf
	} second;
} KDElement;

constant uchar _sceneData[] = {$$SCENEDATA$$};
constant uchar _lightsdata[] = {$$LIGHTSDATA$$};

void setPixel(global uchar4* line, size_t x, uchar r, uchar g, uchar b)
{
	global uchar* pixel = (global uchar*)(line + x);

	pixel[0] = b;	//blue
	pixel[1] = g;	//green
	pixel[2] = r;	//red
	pixel[3] = 255;	//alpha
}

float4 getRayDirection(int width, int height, int x, int y, float4 forward, float4 right, float4 up)
{
	float ratio = (float)width / height;
	float recenteredX = (x - (width/2)) / (2.0f * width) * ratio;
	float recenteredY = (y - (height/2)) / (2.0f * height) ;
	return fast_normalize(forward + (recenteredX * right) + (recenteredY * up));
}

float4 getNormal(constant Object* object, float4 pos)
{
	switch (object->type)
	{
		case 0:		//sphere
			return fast_normalize(pos - object->properties.sphere.position);
		case 1:		//plane
			return fast_normalize(object->properties.plane.normal);
		case 2:		//triangle
			return fast_normalize(cross(object->properties.triangle.p2 - object->properties.triangle.p1, 
										object->properties.triangle.p3 - object->properties.triangle.p1));
	}
	return (float4)(0.0f, 0.0f, 0.0f, 0.0f);
}

bool getRefractionRay(constant Object* object, float4 pos, float4 rayDirection, private float4* outStart, private float4* outDir)
{
	if ((object->surface != 2) && (object->surface != 5))
		return false;

	float4 normal = getNormal(object, pos);
	float4 refractDir;
	float4 VplusN;
	float r;
	float kf;
	float k;

	switch (object->type)
	{
		case 0:		//sphere
			k = 1.01f;		//index of refraction

			r = dot(-normal, rayDirection);
			VplusN = (normal + (rayDirection / r));
			kf = pow(k, 2) * pow(1.0f/r, 2) - dot(VplusN, VplusN);
			if (kf < 0.0f)		//total internal reflection
			{
				return false;
			}
			kf = 1.0f/sqrt(kf);
			refractDir = fast_normalize(-normal + kf * VplusN);
			r = dot(-normal, refractDir);
			*outStart = pos + refractDir * r * (2*object->properties.sphere.radius);

			//*outStart is now the point where the ray comes out of the sphere again.
			//the direction inside is refractDir. However, by leaving the sphere medium, the direction has to be modified again:
			normal = -getNormal(object, *outStart);
			r = dot(-normal, refractDir);
			VplusN = (normal + (refractDir / r));
			kf = pow(1.0f/k, 2) * pow(1.0f/r, 2) - dot(VplusN, VplusN);
			if (kf < 0.0f)		//total internal reflection
			{
				return false;
			}
			kf = 1.0f/sqrt(kf);
			*outDir = fast_normalize(-normal + kf * VplusN);

			break;

		case 1:		//plane
			k = 1.11f;		//index of refraction

			r = dot(-normal, rayDirection);
			VplusN = (normal + (rayDirection / r));
			kf = pow(k, 2) * pow(1.0f/r, 2) - dot(VplusN, VplusN);
			if (kf < 0.0f)		//total internal reflection
			{
				return false;
			}
			kf = 1.0f/sqrt(kf);
			*outDir = fast_normalize(-normal + kf * VplusN);
			*outStart = pos - normal * 0.03;
			break;

		default:
			return false;
	}

	return true;
}

bool intersectBBox(float4 rayStart, float4 rayDirection, float4 bboxMin, float4 bboxMax, float* tmin, float* tmax)
{
	float t0 = FLT_MIN;
	float t1 = FLT_MAX;
	
	float invRayDir = 1.0f / rayDirection.S0;
	float tNear = (bboxMin.S0 - rayStart.S0) * invRayDir;
	float tFar = (bboxMax.S0 - rayStart.S0) * invRayDir;
	
	if (tNear > tFar)
	{
		invRayDir = tFar;
		tFar = tNear;
		tNear = invRayDir;
	}
	t0 = tNear > t0 ? tNear : t0;
	t1 = tFar < t1 ? tFar : t1;
	if (t0 > t1)
		return false;
				
	invRayDir = 1.0f / rayDirection.S1;
	tNear = (bboxMin.S1 - rayStart.S1) * invRayDir;
	tFar = (bboxMax.S1 - rayStart.S1) * invRayDir;
	
	if (tNear > tFar)
	{
		invRayDir = tFar;
		tFar = tNear;
		tNear = invRayDir;
	}
	t0 = tNear > t0 ? tNear : t0;
	t1 = tFar < t1 ? tFar : t1;
	if (t0 > t1)
		return false;

	invRayDir = 1.0f / rayDirection.S2;
	tNear = (bboxMin.S2 - rayStart.S2) * invRayDir;
	tFar = (bboxMax.S2 - rayStart.S2) * invRayDir;
	
	if (tNear > tFar)
	{
		invRayDir = tFar;
		tFar = tNear;
		tNear = invRayDir;
	}
	t0 = tNear > t0 ? tNear : t0;
	t1 = tFar < t1 ? tFar : t1;
	if (t0 > t1)
		return false;

	*tmin = t0;
	*tmax = t1;

	return true;
}

bool intersectingSphere(float4 pos, float radius, float4 rayStart, float4 rayDirection, private float* distanceOut)
{
	float denom;
	float disc;
	float4 eo;

	eo = pos - rayStart;
	denom = dot(eo, rayDirection);
	if (denom >= 0)
	{
		disc = pow(radius, 2) - (dot(eo, eo) - pow(denom, 2));
		if (disc >= 0.0f)
		{
			*distanceOut = denom - sqrt(disc);
			return true;
		}
	}
	return false;
}

bool intersectingTriangle(float4 p1, float4 p2, float4 p3, float4 rayStart, float4 rayDirection, private float* distanceOut)
{
	float4 e1 = p2 - p1;
	float4 e2 = p3 - p1;
	float4 s1 = cross(rayDirection, e2);
	float divisor = dot(s1, e1);
	if (divisor == 0.0f)
		return false;
	float invDivisor = 1.0f / divisor;

	float4 d = rayStart - p1;
	float b1 = dot(d, s1) * invDivisor;
	if ((b1 < 0.0f) || (b1 > 1.0f))
		return false;

	float4 s2 = cross(d, e1);
	float b2 = dot(rayDirection, s2) * invDivisor;
	if ((b2 < 0.0f) || (b1 + b2 > 1.0f))
		return false;

	*distanceOut = dot(e2, s2) * invDivisor;
	return true;
}

//Implements bruteforce rendering (i.e. Testing the ray against every given object):
constant Object* intersectingObjectInNode(float4 rayStart, float4 rayDirection, constant Object* objects, int numObjects, private float* distanceOut)
{
	float minDistance = FLT_MAX;
	constant Object* minObject = 0;
		
	float distance = FLT_MAX;
	float t;

	for (int i = 0; i < numObjects; i++)
	{	
		switch (objects->type)
		{		
			case 0:		//sphere
				intersectingSphere(objects->properties.sphere.position, objects->properties.sphere.radius, rayStart, rayDirection, &distance);
				break;
			case 1:		//plane
				t = dot(fast_normalize(objects->properties.plane.normal), rayDirection);
				if (t <= 0.0f)
				{
					distance = (dot(fast_normalize(objects->properties.plane.normal), rayStart) + objects->properties.plane.offset) / (-t);					
				}
				break;
			case 2:		//triangle
				intersectingTriangle(objects->properties.triangle.p1, objects->properties.triangle.p2, objects->properties.triangle.p3, rayStart, rayDirection, &distance);								
				break;
		}
		
		if (distance < minDistance && distance > 0)
		{
			minObject = objects;
			minDistance = distance;
		}

		objects++;
	}
	if (minDistance == FLT_MAX)
		return 0;

	*distanceOut = minDistance;
	return minObject;
}

//Implements KD-Tree algorithm:
constant Object* intersectingObject(float4 rayStart, float4 rayDirection, private float* distanceOut)
{
	constant uchar* sceneData = _sceneData;
	int objPointer = *((constant int*)sceneData);
	constant Object* objects = (constant Object*)(sceneData + objPointer);
	constant Object* intersectingObject = 0;
	int numPlanes = *((constant int*)(sceneData+4));
	int numObjects = *((constant int*)(sceneData+8));

#ifdef  _BRUTEFORCERENDERING_
	return intersectingObjectInNode(rayStart, rayDirection, objects, numObjects, distanceOut);
#endif

	float planeDistanceOut;
	constant Object* intersectingPlane = intersectingObjectInNode(rayStart, rayDirection, objects, numPlanes, &planeDistanceOut);
	
	sceneData += 4*4;
	float4 bboxMin = *((constant float4*)sceneData);
	sceneData += 4*4;
	float4 bboxMax = *((constant float4*)sceneData);
	sceneData += 4*4;
	constant KDElement* kdroot = (constant KDElement*)sceneData;
	constant KDElement* kdcurrent = kdroot;

#if _SHORT_STACK_SIZE > 0
	constant KDElement* shortStackKDElements[_SHORT_STACK_SIZE];	
	private float shortStackTMin[_SHORT_STACK_SIZE];
	private float shortStackTMax[_SHORT_STACK_SIZE];
	uchar stackPos = 0;
	uchar stackCount = 0;
#endif

	float global_tmin;
	float global_tmax;
	float tmin;
	float tmax;
	float thit;

	uint first, second;

	float rayAxisDirection, rayAxisStart;
	
	if (!intersectBBox(rayStart, rayDirection, bboxMin, bboxMax, &global_tmin, &global_tmax))
	{
		*distanceOut = planeDistanceOut;
		return intersectingPlane;
	}

	tmin = global_tmin;
	tmax = global_tmax;

	do
	{
		if (((kdcurrent->first.flags) & 3) == 0)		//is leaf
		{
			if (kdcurrent->first.numObjects != 0)
			{
				intersectingObject = intersectingObjectInNode(rayStart, rayDirection, 
															  objects + kdcurrent->second.objectIndex, 
															  kdcurrent->first.numObjects >> 2, distanceOut);				
			}
			if (intersectingObject == 0)
			{
#if _SHORT_STACK_SIZE > 0
				if (stackCount > 0)		//use shortstack
				{
					kdcurrent = shortStackKDElements[stackPos];
					tmin = shortStackTMin[stackPos];
					tmax = shortStackTMax[stackPos];
					if (stackPos > 0)
					{
						stackPos--;
					}
					else
					{
						stackPos = _SHORT_STACK_SIZE-1;
					}
					stackCount--;
					continue;
				}
#endif

				//use "kd restart" mechanism:				
				if (tmax >= global_tmax)
					break;
				//continue:
				tmin = tmax;
				tmax = global_tmax;
				kdcurrent = kdroot;
			}
		}
		else										//is node
		{
			switch ((kdcurrent->first.flags) & 3)
			{
				case 1:	//X Axis
					rayAxisDirection = rayDirection.S0;
					rayAxisStart = rayStart.S0;
					break;
				case 2:	//Y Axis
					rayAxisDirection = rayDirection.S1;
					rayAxisStart = rayStart.S1;
					break;
				case 3:	//Z Axis
					rayAxisDirection = rayDirection.S2;
					rayAxisStart = rayStart.S2;
					break;
			}

			thit = ((kdcurrent->first.split) - rayAxisStart) / rayAxisDirection;

			if (rayAxisStart < (kdcurrent->first.split))
			{
				first = 1;
				second = kdcurrent->second.aboveChild;
			}
			else
			{
				first = kdcurrent->second.aboveChild;
				second = 1;
			}

			if ((thit >= tmax) || (thit < 0))
			{
				kdcurrent = kdcurrent + first;
			}
			else if (thit <= tmin)
			{
				kdcurrent = kdcurrent + second;
			}
			else
			{
#if _SHORT_STACK_SIZE > 0
				//push second node on stack:
				stackPos = (stackPos + 1) % _SHORT_STACK_SIZE;
				if (stackCount < _SHORT_STACK_SIZE)
					stackCount++;
				shortStackKDElements[stackPos] = kdcurrent + second;
				shortStackTMin[stackPos] = thit;
				shortStackTMax[stackPos] = tmax;
#endif

				//current node is the first one:
				tmax = thit;
				kdcurrent = kdcurrent + first;
			}
		}
	} while (intersectingObject == 0);

	if ((intersectingObject == 0) || (intersectingPlane != 0 && (*distanceOut > planeDistanceOut)))
	{
		intersectingObject = intersectingPlane;
		*distanceOut = planeDistanceOut;
	}

	return intersectingObject;
}

/* used for testing the shadow ray: */
float testRay(float4 rayStart, float4 rayDirection)
{
	float distance;
	if (intersectingObject(rayStart, rayDirection, &distance) == 0)
		return 0.0f;

	return distance;
}

float getRoughness(uchar surface)
{
	switch (surface)
	{
		case 0:	//checkerboard
			return 150.0f;
		case 1: //shiny
			return 50.0f;
		case 2: //transparent
		case 5:	//transparent and shiny
			return 10.0f;
		case 3:	//sky
		case 4:	//wood
			return 0.0f;
	}
	return 0.0f;
}

float getReflect(uchar surface, float4 pos)
{
	switch (surface)
	{
		case 0:	//checkerboard
			return (int)(floor(pos.s2) + floor(pos.s0)) % 2 != 0 ? 0.05f : 0.1f;
		case 1: //shiny
			return 0.7f;
		case 2: //transparent
		case 5:	//transparent and shiny
			return -0.9f;
		case 3:	//sky
		case 4:	//wood
			return 0.0f;
	}
	return 0.0f;
}

float4 getSpecular(uchar surface, float4 pos)
{
	switch (surface)
	{
		case 0:	//checkerboard
			return (float4)(1.0f, 1.0f, 1.0f, 0.0f);
		case 1: //shiny
			return (float4)(0.5f, 0.5f, 0.5f, 0.0f);
		case 2: //transparent
		case 5:	//transparent and shiny
			return (float4)(0.5f, 0.5f, 0.5f, 0.0f);
		case 3:	//sky
		case 4:	//wood
			return (float4)(0.0f, 0.0f, 0.0f, 0.0f);
	}
	return (float4)(0.0f, 0.0f, 0.0f, 0.0f);
}

float4 getDiffuse(constant Object* object, float4 pos, read_only image2d_t skyTexture, read_only image2d_t woodTexture, sampler_t sampler)
{
	uint4 rgba;
	float u, v;

	switch (object->surface)
	{
		case 0:	//checkerboard
			return (int)(floor(pos.s2) + floor(pos.s0)) % 2 != 0 ? (float4)(1.0f,1.0f,1.0f,0.0f) : (float4)(0.0f,0.0f,0.0f,0.0f);
		case 1: //shiny
			return (float4)(1.0f, 1.0f, 1.0f, 0.0f);
		case 2: //transparent
		case 5:	//transparent and shiny
			return (float4)(1.0f, 1.0f, 1.0f, 0.0f);
		case 3:	//sky
		case 4:	//wood
			v = sqrt(pow(pos.s0, 2) + pow(pos.s1, 2) + pow(pos.s2, 2));
			u = pos.s0 / v;
			v = pos.s1 / v;
			if (object->surface == 3)
				rgba = read_imageui(skyTexture, sampler, (float2)(u, v));
			else
				rgba = read_imageui(woodTexture, sampler, (float2)(u, v));
			return (float4)(rgba.s2 / 255.0f, rgba.s1 / 255.0f, rgba.s0 / 255.0f, 0.0f);		
	}
	return (float4)(0.0f, 0.0f, 0.0f, 0.0f);
}

float4 getNaturalColor(float4 pos, float4 normal, float4 rd, constant Object* object, read_only image2d_t skyTexture, read_only image2d_t woodTexture, sampler_t sampler)
{
	int numLights = *((constant int*)_lightsdata);
	constant Light* lights = (constant Light*)(_lightsdata + 4*4);
	bool isInShadow = false;

	float4 ret = (float4)(0.0f,0.0f,0.0f,0.0f);

	for (int i = 0; i < numLights; i++)
	{
		float4 ldis = lights[i].position - pos;
		float4 livec = fast_normalize(ldis);

#ifdef _SHOW_SHADOWS_
		if (object->surface != 3 && object->surface != 2 && object->surface != 5)	//sky and transparent don't receive shadows
		{			
			float neatIsect = testRay(pos, livec);
			isInShadow = !((neatIsect > length(ldis)) || (neatIsect == 0.0f));
		}
#endif
		
		if (!isInShadow)
        {
            float illum = dot(livec, normal);
			float4 lcolor;
			if(illum > 0){
				lcolor= illum * lights[i].color;
			}else{
				lcolor=(float4)(0.0f,0.0f,0.0f,0.0f);
			}

			float roughness = getRoughness(object->surface);

			float spec = dot(livec, fast_normalize(rd));
			float4 scolor;
			if (spec > 0.0f) {
				scolor = pow(spec, roughness) * lights[i].color;
			}else{
				scolor = (float4)(0.0f,0.0f,0.0f,0.0f);
			}
			
			float4 diffuse = getDiffuse(object, pos, skyTexture, woodTexture, sampler);
			float4 specular = getSpecular(object->surface, pos);

            ret = ret + (float4)(lcolor.s0*diffuse.s0, lcolor.s1*diffuse.s1, lcolor.s2*diffuse.s2, 0.0f) 
					  + (float4)(scolor.s0*specular.s0, scolor.s1*specular.s1, scolor.s2*specular.s2, 0.0f);
        }
	}

	return ret;
}

float4 traceRay(float4 rayStart, float4 rayDirection, read_only image2d_t skyTexture, read_only image2d_t woodTexture, sampler_t sampler)
{
	float4 color = (float4)(0.0f,0.0f,0.0f,0.0f);
	float reflect = 1.0f;
	int depth = _DEPTH;

	constant Object* object2;
	float4 pos2;
	float4 normal2;
	float4 reflectDir2;

	for (int i = 0; i < depth; i++)
	{
		float distance;
		constant Object* object = intersectingObject(rayStart, rayDirection, &distance);
		if (object == 0)
			return color;

		float4 pos = (distance * rayDirection) + rayStart;

		//shade:
		float4 normal = getNormal(object, pos);
		float4 reflectDir = rayDirection - ((2 * dot(normal, rayDirection)) * normal);
				
		color = color + reflect * getNaturalColor(pos, normal, reflectDir, object, skyTexture, woodTexture, sampler);

		//reflection/refraction:
		reflect *= getReflect(object->surface, pos);
		if (reflect == 0.0f)
			break;
		if (reflect < 0.0f)
		{
			reflect = -reflect;

			if ((i < depth-1) && (object->surface == 5))	//for surface "transparent and shiny"
			{
				//compute the color of the reflection ray with depth 1 (to fake the ability to reflect and refract at the same time):
				object2 = intersectingObject(pos, reflectDir, &distance);			
				if (object2 != 0)
				{
					pos2 = (distance * reflectDir) + pos;
					normal2 = getNormal(object2, pos2);
					reflectDir2 = reflectDir - ((2 * dot(normal2, reflectDir)) * normal2);
					color = color + reflect * getNaturalColor(pos2, normal2, reflectDir2, object2, skyTexture, woodTexture, sampler);
				}
			}

			if (getRefractionRay(object, pos, rayDirection, &rayStart, &rayDirection))
				continue;
		}
		rayStart = pos;
		rayDirection = reflectDir;
	}
	return color;
}

kernel void RayTrace (int width, int height, global uchar* pOutput, int outputStride, float4 cameraPosition, float4 cameraForward, float4 cameraRight, float4 cameraUp, read_only image2d_t skyTexture, read_only image2d_t woodTexture, sampler_t sampler)
{
	size_t x = get_global_id(0);
	size_t y = get_global_id(1);	
	global uchar4* pO = (global uchar4*)(pOutput+y*outputStride);
	
	float4 rayDirection = getRayDirection(width, height, x, y, cameraForward, cameraRight, cameraUp);
	float4 pixelColor = traceRay(cameraPosition, rayDirection, skyTexture, woodTexture, sampler);	
	setPixel(pO, x, (int)(pixelColor.s0 > 1 ? 255 : pixelColor.s0 * 255), 
					(int)(pixelColor.s1 > 1 ? 255 : pixelColor.s1 * 255), 
					(int)(pixelColor.s2 > 1 ? 255 : pixelColor.s2 * 255));
}
